from pathlib import Path
import numpy as np
from collections import defaultdict
from tqdm import tqdm
import matplotlib.pyplot as plt

from utils import load_record


DATASETS = {
    'mit-bih': Path('processed_data/mit-bih'),
    'mit-bih-supra': Path('processed_data/mit-bih-supra'),
    'incart': Path('processed_data/incart')
}

if __name__ == '__main__':
    total_counts = defaultdict(int)
    for n, p in DATASETS.items():
        counts = defaultdict(int)
        img_gen = False
        for patient_path in tqdm(list(p.glob('PATIENT_*'))):
            for record_path in patient_path.glob('RECORD_*'):
                X, X_cwt, X_rr, y, loc = load_record(record_path)
                labels, label_counts = np.unique(y, return_counts=True)
                for l, c in zip(labels, label_counts):
                    counts[l] += c
        print(n)
        for l, c in counts.items():
            print(f"{l}: {c}")
            total_counts[l] += c
        print()
    print("TOTAL")
    for i in range(len(total_counts)):
        print(f"{i}: {total_counts[i]}")